import rosbag
import rospy
from unitree_legged_msgs.msg import LowState

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, default='')
parser.add_argument('-o', '--output', type=str, default='')

args = parser.parse_args()
input_bag = args.input
output_bag = args.output

with rosbag.Bag(output_bag, 'w') as outbag:
    for topic, msg, t in rosbag.Bag(input_bag).read_messages():
        if topic == '/state_estimator/log':
            d = list(msg.data)
            d[110] = d[110] + d[119]
            d[111] = d[111] + d[120]
            d[112] = d[112] + d[121]
            d[113] = d[113] + d[116]
            d[114] = d[114] + d[117]
            d[115] = d[115] + d[118]
            msg.data = tuple(d)
        outbag.write(topic, msg, msg.header.stamp if msg._has_header else t)
